import os, sys
sys.path.append(".")

from unifuncs import *
import pyarma as pa
import numpy as np

# 真解

def Maxwell(M, lb, ub, delta = 0.04):

    h = (ub - lb) / M
    x = pa.linspace(lb + h, ub - h, M - 1)
    u = pa.zeros(M, 1)
    u[0] = 0
    u[1:M-1, 0] = pa.exp(-pa.pow(x - 7 * pa.ones(M - 1, 1), 2))
    return u

def SWave1(M, lb, ub, t, c = 0.1, delta = 1.0, x0 = 0):

    h = (ub - lb) / M

    p = np.sqrt(c / (4 * delta * (c + 1)))
    v = 1 + c

    u = pa.mat(M, 1)
    u[0] = t
    for i in range(1, M):
        x = lb + i * h
        u[i] = 3 * c / (np.cosh(p * (x - v * t - x0)) ** 2)

    return u

def SWave1Combination(M, lb, ub, k = [0.4, 0.3], delta = 1.0, x0 = [15, 35]):
    
    u = pa.zeros(M, 1)
    for i in range(len(k)):
        c = 4 * (k[i] ** 2) / (1 - 4 * (k[i] ** 2))
        u += SWave1(M, lb, ub, 0, c, delta, x0[i])
    return u

def SWave2(M, lb, ub, t, a = 0.5, beta = 0):

    h = (ub - lb) / M

    A = 3 * (a ** 2) / (1 - (a ** 2))
    k = a / 2
    alpha = a / (2 * (1 - (a ** 2)))

    u = pa.mat(M, 1)
    u[0] = t
    for i in range(1, M):
        x = lb + i * h
        u[i] = A / (np.cosh(k * x - alpha * t + beta) ** 2)

    return u

# 有限差分格式

def FDM_F(u0, M, lb, ub, delta = 1.0):
    
    h = (ub - lb) / M
    un = pa.zeros(pa.size(u0))
    un[0] = 1
    u = u0[1:u0.n_elem-1, 0]
    un[1:u0.n_elem-1, 0] = - (pa.eye(M - 1, M - 1) - delta * diff.cd2(M - 1, h)).i() * (pa.eye(M-1, M-1) + pa.diagmat(u)) * diff.cd1(M - 1, h) * u
    return un

def FDM_J(u0, M, lb, ub, delta = 1.0):
    
    h = (ub - lb) / M
    Jn = pa.zeros(M, M)
    u = u0[1:u0.n_elem-1, 0]
    Jn[1:u0.n_elem-1, 1:u0.n_elem-1] = - (pa.eye(M - 1, M - 1) - delta * diff.cd2(M - 1, h)).i() * ((pa.eye(M-1, M-1) + pa.diagmat(u)) * diff.cd1(M - 1, h) + diff.cd1(M - 1, h) * u)
    return Jn

def SHFDM_F(u0, M, lb, ub, delta = 1.0):

    h = (ub - lb) / M
    un = pa.zeros(pa.size(u0))
    un[0] = 1
    u = u0[1:u0.n_elem-1, 0]
    un[1:u0.n_elem-1, 0] = - (pa.eye(M - 1, M - 1) - delta * diff.cd2(M - 1, h)).i() * diff.cd1(M - 1, h) * (u + pa.pow(u, 2) / 2)
    return un

def SHFDM_J(u0, M, lb, ub, delta = 1.0):

    h = (ub - lb) / M
    Jn = pa.zeros(M, M)
    u = u0[1:u0.n_elem-1, 0]
    Jn[1:u0.n_elem-1, 1:u0.n_elem-1] = - (pa.eye(M - 1, M - 1) - delta * diff.cd2(M - 1, h)).i() * diff.cd1(M - 1, h) * (pa.eye(M - 1, M - 1) + pa.diagmat(u))
    return Jn

def Invariant(c, delta):

    p = np.sqrt(c / (4 * delta * (c + 1)))
    I1 = 6 * c / p
    I2 = 12 * (c ** 2) / p + 48 * p * (c ** 2) * delta / 5
    I3 = 36 * (c ** 2) * (1 + 4 * c / 5) / p
    return [I1, I2, I3]

# Galerkin有限元法

def Galerkin_A(M, lb, ub, delta = 1.0):
    
    h = (ub - lb) / M
    An = pa.zeros(M, M)
    An[1:M-1, 1:M-1] = (diff.GLK1010(M - 1, h) + delta * diff.GLK1111(M - 1, h)).i() * diff.GLK1011(M - 1, h)
    return An

def Galerkin_g(u0, M, lb, ub, delta = 1.0):

    h = (ub - lb) / M
    un = pa.zeros(pa.size(u0))
    un[0] = 1
    u = u0[1:u0.n_elem-1, 0]
    un[1:u0.n_elem-1, 0] = (diff.GLK1010(M - 1, h) + delta * diff.GLK1111(M - 1, h)).i() * diff.GLK2011(M - 1, h, u) / 2
    return un

def Galerkin_F(u0, M, lb, ub, delta = 1.0):

    h = (ub - lb) / M
    un = pa.zeros(pa.size(u0))
    un[0] = 1
    u = u0[1:u0.n_elem-1, 0]
    un[1:u0.n_elem-1, 0] = (diff.GLK1010(M - 1, h) + delta * diff.GLK1111(M - 1, h)).i() * (diff.GLK1011(M - 1, h) * u + diff.GLK2011(M - 1, h, u) / 2)
    return un

def Galerkin_J(u0, M, lb, ub, delta = 1.0):

    h = (ub - lb) / M
    Jn = pa.zeros(M, M)
    for i in range(1, M):

        if i > 1:

            Jn[i, i-1] += 2 * u0[i-1] + u0[i]
            Jn[i, i] += u0[i-1]

        if i < M-1:

            Jn[i, i+1] -= 2 * u0[i+1] + u0[i]
            Jn[i, i] -= u0[i+1]

    Jn /= 6
    Jn[1:M-1, 1:M-1] = (diff.GLK1010(M - 1, h) + delta * diff.GLK1111(M - 1, h)).i() * Jn[1:M-1, 1:M-1]
    Jn += Galerkin_A(M, lb, ub, delta)

    return Jn

# 组合

def Finite_Difference(M, lb, ub, delta):
    
    h = (ub - lb) / M
    F = lambda u: FDM_F(u, M, lb, ub, delta)
    J = lambda u: FDM_J(u, M, lb, ub, delta)
    A = pa.zeros(M, M)
    A[1:M-1, 1:M-1] = - (pa.eye(M - 1, M - 1) - delta * diff.cd2(M - 1, h)).i() * diff.cd1(M - 1, h)
    g = lambda u: F(u) - A * u

    return F, J, g, A

def Invariant_Finite_Difference(M, lb, ub, delta):
    
    h = (ub - lb) / M
    F = lambda u: SHFDM_F(u, M, lb, ub, delta)
    J = lambda u: SHFDM_J(u, M, lb, ub, delta)
    A = pa.zeros(M, M)
    A[1:M-1, 1:M-1] = - (pa.eye(M - 1, M - 1) - delta * diff.cd2(M - 1, h)).i() * diff.cd1(M - 1, h)
    g = lambda u: F(u) - A * u

    return F, J, g, A

def Invariant_Finite_Element(M, lb, ub, delta):

    F = lambda u: Galerkin_F(u, M, lb, ub, delta)
    g = lambda u: Galerkin_g(u, M, lb, ub, delta)
    J = lambda u: Galerkin_J(u, M, lb, ub, delta)
    A = Galerkin_A(M, lb, ub, delta)

    return F, J, g, A

def get_Errors(M, lb, ub, u, ue):

    h = (ub - lb) / M
    L1 = h * pa.norm(pa.abs(u - ue(u[0])), 1)
    L2 = np.sqrt(h) * pa.norm(pa.abs(u - ue(u[0])), 2)
    Li = pa.norm(pa.abs(u - ue(u[0])), 'inf')

    return [L1, L2, Li]

def get_Invariant(M, lb, ub, u, delta = 1.0):

    h = (ub - lb) / M
    Ms = h * pa.sum(u[1:M-1, 0])[0]
    Ps = h * pa.sum(pa.pow(u[1:M-1, 0], 2) + delta * pa.pow(diff.cd1(M - 1, h) * u[1:M-1, 0], 2))[0]
    Es = h * pa.sum(pa.pow(u[1:M-1, 0], 3) + 3 * pa.pow(u[1:M-1, 0], 2))[0]

    return [Ms, Ps, Es]